package metrics

import (
	"strings"
	"time"

	"gitee.com/zawei/yggdrasil/pkg/lib"
	"gitee.com/zawei/yggdrasil/pkg/lib/gorm/plugin"
	"gitee.com/zawei/yggdrasil/pkg/metrics"
	"gorm.io/gorm"
)

func init() {

}

type metric struct {
	addr string
}

func New(addr string) gorm.Plugin {
	return &metric{addr: addr}
}

func (m *metric) Name() string {
	return "gorm:metric"
}

func (m *metric) Initialize(db *gorm.DB) error {
	m.registerCallbacks(db)
	return nil
}

func (m *metric) statistic(next func(*gorm.DB)) func(*gorm.DB) {
	return func(db *gorm.DB) {
		beg := time.Now()
		next(db)
		cost := time.Since(beg)

		statement := plugin.StatementToString(db.Statement, false)
		method := db.Statement.Schema.Name + "." + db.Statement.Table + "." + (strings.Split(statement, " "))[0]
		// error metric
		if db.Error != nil && db.Error != gorm.ErrRecordNotFound {
			_ = metrics.CounterAdd(lib.MetricsLibHandleTotal, 1, map[string]string{
				"type": "grom", "method": method, "address": m.addr, "code": "ERR"})
		} else {
			_ = metrics.CounterAdd(lib.MetricsLibHandleTotal, 1, map[string]string{
				"type": "grom", "method": method, "address": m.addr, "code": "OK"})
		}
		_ = metrics.HistogramObserve(lib.MetricsLibHandleSec, cost.Seconds(), map[string]string{
			"type": "grom", "method": method, "address": m.addr})
	}
}

func (m *metric) registerCallbacks(db *gorm.DB) {
	_ = db.Callback().Create().Replace("gorm:create", m.statistic(db.Callback().Create().Get("gorm:create")))
	_ = db.Callback().Query().Replace("gorm:query", m.statistic(db.Callback().Query().Get("gorm:query")))
	_ = db.Callback().Delete().Replace("gorm:delete", m.statistic(db.Callback().Delete().Get("gorm:delete")))
	_ = db.Callback().Update().Replace("gorm:update", m.statistic(db.Callback().Update().Get("gorm:update")))
	_ = db.Callback().Row().Replace("gorm:row", m.statistic(db.Callback().Row().Get("gorm:row")))
	_ = db.Callback().Raw().Replace("gorm:raw", m.statistic(db.Callback().Raw().Get("gorm:raw")))
}
